import logging
import re
import json
import xml.etree.ElementTree as ET
import asyncio
from typing import Optional, List, Dict, Any
import traceback
from maas.constants.http_method_name import HttpMethodName
from maas.http.api_client import ApiClient
from maas.model.api_request import ApiRequest
from configparser import ConfigParser
import requests
import time

EXT_API_URL = None
EXT_API_KEY = None
EXT_API_SECRET = None
EXT_MODEL_NAME = None

# --- 日志配置 ---
logging.basicConfig(format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
                    datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO)
logger = logging.getLogger(__name__)


# --- 配置初始化函数  ---
def init_ext(config: ConfigParser):
    global EXT_API_URL, EXT_API_KEY, EXT_API_SECRET, EXT_MODEL_NAME
    try:
        section = 'ext_model'
        EXT_API_URL = config.get(section, 'api_url')
        EXT_API_KEY = config.get(section, 'api_key')
        EXT_API_SECRET = config.get(section, 'api_secret')
        EXT_MODEL_NAME = config.get(section, 'model_name')
        logger.info(f"EXT服务配置加载成功, 模型: {EXT_MODEL_NAME}")
    except Exception as e:
        logger.error(f"加载EXT服务配置失败: {e}")
        EXT_API_URL, EXT_API_KEY, EXT_API_SECRET, EXT_MODEL_NAME = None, None, None, None


# --- 核心API调用函数 ---
def call_llm_api(messages: List[Dict], max_tokens: int = 6000, temperature: float = 0.1) -> Optional[str]:
    if not all([EXT_API_URL, EXT_API_KEY, EXT_API_SECRET, EXT_MODEL_NAME]):
        logger.error("EXT服务配置不完整，无法调用LLM API。")
        return None
    api_request = ApiRequest(HttpMethodName.POST, EXT_API_URL)
    api_request.add_headers("Content-Type", "application/json")
    api_request.set_credential(EXT_API_KEY, EXT_API_SECRET)
    request_body = {
        "model": EXT_MODEL_NAME,
        "messages": messages,
        "max_tokens": max_tokens,
        "temperature": temperature,
        "stream": False
    }
    api_request.set_json_body(json.dumps(request_body))
    try:
        response = ApiClient().send_request(api_request)
        if response.ok:
            response_data = response.json()
            return response_data.get("choices", [{}])[0].get("message", {}).get("content", "")
        else:
            logger.error(f"聊天API请求失败，状态码 {response.status_code}: {response.text}")
            return None
    except Exception:
        logger.error(f"调用聊天API时发生异常: {traceback.format_exc()}")
        return None


# --- 长文本处理模块  ---
def split_text(text: str, chunk_size: int = 25000, overlap: int = 500) -> List[str]:
    if len(text) <= chunk_size: return [text]
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end])
        if end >= len(text): break
        start = end - overlap
    logger.info(f"文本被分割为 {len(chunks)} 个片段进行并行处理。")
    return chunks


def extract_from_chunk_sync(chunk: str, instruction: str, output_format: str, chunk_number: int, total_chunks: int) -> \
Dict[str, str]:
    task_start_time = time.time()
    logger.info(f"【线程启动】抽取任务 {chunk_number}/{total_chunks}...")
    prompt_str = generate_prompt(chunk, instruction, output_format)
    messages = parse_prompt_to_messages(prompt_str)
    model_response = call_llm_api(messages)
    if model_response is None:
        logger.error(f"任务 {chunk_number}/{total_chunks} 模型API调用失败。")
        return {}
    result_data = parse_output(model_response)
    task_end_time = time.time()
    duration = task_end_time - task_start_time
    logger.info(f"【线程结束】抽取任务 {chunk_number}/{total_chunks}，耗时: {duration:.2f} 秒。")
    return result_data


async def final_extraction_reduce(raw_results: List[Dict], instruction: str, output_format: str) -> Dict[str, str]:
    logger.info("开始对所有分片结果进行最终整合...")
    context_for_reduce = ""
    for i, result_dict in enumerate(raw_results):
        context_for_reduce += f"--- 来自片段 {i + 1} 的初步抽取结果 ---\n"
        for key, value in result_dict.items():
            context_for_reduce += f"{key}: {value}\n"
        context_for_reduce += "\n"
    reduce_prompt_str = f"""<|im_start|>/no_think system
你是一个顶级的、专业的数据整合与清洗专家。你的任务是分析下面提供的、从多个文本片段中初步抽取的、可能包含重复或无效信息（如 "无"）的数据集合。请根据用户的最终提取要求，对这些数据进行去重、整合和清洗，并生成一份唯一的、准确的最终结果。
要求：
1.  **去重**：如果多个片段都提到了相同的信息（例如同一个参会人员），在最终结果中只保留一次。不同的片段提到相同问题的答案注意两者都要输出，并以多个参数分开回复，相同出参也无所谓。
2.  **整合**：将不同片段抽取的同类信息（例如多个"关键词"）合并到一起，并用分号(;)分隔。
3.  **清洗**：忽略所有值为 "无" 或明显无效的条目，除非所有片段对某一字段的提取结果都是 "无"。
4.  **严格遵循格式**：最终输出必须严格遵守用户的格式要求（JSON或Markdown）。
用户的最终提取要求是：
---
{instruction}
---
<|im_start|>user
待整合的初步抽取数据集合如下：
---
{context_for_reduce}
---
请根据上述“用户的最终提取要求”，生成最终的、干净的提取结果。<|im_end|>
<|im_start|>assistant
"""
    final_prompt_template = generate_prompt("", instruction, output_format)
    final_prompt_str = re.sub(r'<\|im_start\|>user.*<\|im_start\|>assistant',
                              reduce_prompt_str.split('<|im_start|>user')[1], final_prompt_template, flags=re.DOTALL)
    messages = parse_prompt_to_messages(final_prompt_str)
    loop = asyncio.get_event_loop()
    final_response = await loop.run_in_executor(None, call_llm_api, messages)
    if final_response is None:
        raise ValueError("最终整合阶段的模型API调用失败。")
    final_data = parse_output(final_response)
    logger.info("最终结果整合完成。")
    return final_data


async def process_long_text_map_reduce(text: str, instruction: str, output_format: str) -> Dict[str, str]:
    chunks = split_text(text)
    logger.info(f"开始并行派发 {len(chunks)} 个抽取任务...")
    map_start_time = time.time()
    tasks = [asyncio.to_thread(extract_from_chunk_sync, chunk, instruction, output_format, i + 1, len(chunks)) for
             i, chunk in enumerate(chunks)]
    results_from_chunks = await asyncio.gather(*tasks)
    map_end_time = time.time()
    logger.info(f"所有 {len(chunks)} 个并行抽取任务完成，总耗时: {map_end_time - map_start_time:.2f} 秒。")
    final_result = await final_extraction_reduce(results_from_chunks, instruction, output_format)
    return final_result


# --- 功能函数  ---
def generate_prompt(text: str, instruction: str, output_format: str) -> str:
    processed_text = text
    system_base = """<|im_start|>/no_think system
你是一个专业的信息提取助手。请认真分析用户要求，分析用户的抽取点，从提供的文档内容中提取信息。
要求：
1. 严格按照用户要求提取信息，并列关系的抽取点分开处理。
2. 保证信息绝对准确且具体存在，不存在的内容写"无"。
3. 文本内容中去掉多余的空格和换行符。
4. 确保提取的内容格式规范化，特别是英文、数字和标点符号，英文注意使用半角而非全角。
5. 严格从原文中直接提取内容，不要对提取出的结果进行任何形式的翻译。提取结果的语言应与原文保持一致。"""
    if output_format.lower() == "markdown":
        system_prompt = f"""... (markdown prompt 保持不变) ..."""
    else:
        system_prompt = f"""
{system_base}
6. 返回严格的扁平化JSON格式（单层结构，禁止任何嵌套），使用中文键名，确保语法正确。
7. 只输出纯JSON对象，不要包含任何解释或代码块标记 (```)。
8. 如果用户要求模糊（如"识别文中要素"），请根据上下文尽可能全面提取关键信息，包括但不限于日期、人物、公司、职位、事件等。
9. JSON格式规范：
    - 必须是单层键值对结构。
    - 值只能是字符串类型（禁止数组、对象等嵌套结构）。
    - **多个值**必须用**分号(;)**分隔在同一个字符串中。
    - **示例格式**：{{"会议时间": "20XX-XX-XX XX:XX", "主持人": "XX", "参会人员": "张三;李四;王五", "关键词": "关键词1;关键词2;关键词3", "摘要": "XXX"}}
10. 特别注意：
    - 使用半角符号：冒号(:)、逗号(,)、分号(;)、引号(")。
    - 键名使用双引号包裹。
    - 字符串值使用双引号包裹。
    - **键值对之间**用**逗号**分隔。
    - 整个JSON对象不要换行和缩进（输出在一行内）。
    - 确保JSON字符串语法正确，JSON字符串值中如果包含特殊字符，如反斜杠(\)或引号(")，必须严格按照JSON标准进行转义。例如，一个单独的反斜杠应表示为'\\'。
用户的具体要求（提取字段）：
{instruction}
<|im_start|>user
文档内容：
{processed_text}<|im_end|>
<|im_start|>assistant
"""
    return system_prompt


def parse_prompt_to_messages(prompt: str) -> List[Dict[str, str]]:
    messages = []
    parts = re.split(r'(<\|im_start\|>|<\|im_end\|>)', prompt)
    role, content = None, ""
    for part in parts:
        part = part.strip()
        if not part: continue
        if part == '<|im_start|>':
            if role and content: messages.append({'role': role, 'content': content.strip()})
            content = ""
        elif part == '<|im_end|>':
            if role and content: messages.append({'role': role, 'content': content.strip()})
            role, content = None, ""
        elif part.startswith('/no_think system'):
            role, content = 'system', part.replace('/no_think system', '').strip()
        elif part == 'user':
            role, content = 'user', ""
        elif part == 'assistant':
            role, content = 'assistant', ""
        else:
            if role: content += part + " "
    if role and content.strip(): messages.append({'role': role, 'content': content.strip()})
    if len(messages) > 1:
        merged_messages = [messages[0]]
        for i in range(1, len(messages)):
            if messages[i]['role'] == merged_messages[-1]['role']:
                merged_messages[-1]['content'] += "\n" + messages[i]['content']
            else:
                merged_messages.append(messages[i])
        return merged_messages
    return messages


# --- 主处理函数  ---
async def process_content(text: str, instruction: str, blocks_info: Optional[List[Dict]] = None,
                          output_format: str = "json") -> Dict[str, Any]:
    MAX_LENGTH = 25000
    try:
        result_data = {}
        if len(text) > MAX_LENGTH:
            logger.info(f"检测到长文本 (长度: {len(text)} > {MAX_LENGTH})，启动两阶段 Map-Reduce 抽取流程...")
            result_data = await process_long_text_map_reduce(text, instruction, output_format)
        else:
            logger.info(f"文本长度为 {len(text)}，使用标准处理流程...")
            loop = asyncio.get_event_loop()
            prompt_str = generate_prompt(text, instruction, output_format)
            messages = parse_prompt_to_messages(prompt_str)
            model_response = await loop.run_in_executor(None, call_llm_api, messages)
            if model_response is None: raise ValueError("模型API调用失败，未返回任何内容。")
            result_data = parse_output(model_response)
        formatted_data = format_output(result_data, output_format)
        return {"data": {"extracted": result_data, "formatted_data": formatted_data},
                "is_supported_format": True, "message": "处理成功", "code": 200}
    except Exception as e:
        logger.error(f"模型处理失败: {str(e)}", exc_info=True)
        return {"data": {"extracted": {}, "formatted_data": None},
                "is_supported_format": True, "message": f"信息提取失败: {str(e)}", "code": 500}


# --- 增强的输出解析函数 ---
def parse_output(output: str) -> Dict[str, str]:
    """
    增强版的解析函数，能处理多种不规范的JSON输出。
    """
    if not isinstance(output, str):
        return {}

    output = output.strip()

    # 步骤 1: 尝试从 markdown 代码块中提取
    json_match = re.search(r'```json\s*(\{.*?})\s*```', output, re.DOTALL)
    if json_match:
        output = json_match.group(1).strip()

    # 步骤 2: 尝试处理被额外引号包裹的JSON字符串，例如 "{\"key\": \"value\"}"
    if output.startswith('"') and output.endswith('"'):
        try:
            # json.loads 可以将这种带引号的字符串正确解码为内部的字符串
            output = json.loads(output)
        except json.JSONDecodeError:
            # 如果解码失败，说明不是标准的带引号字符串，继续使用原始字符串
            pass

    if isinstance(output, str):
        output = re.sub(r'\\(?!["\\/bfnrtu])', r'\\\\', output)

    # 步骤 3: 核心解析逻辑
    try:
        # 此时 output 可能是干净的字符串或上一步解码后的字典
        if isinstance(output, str) and output.startswith('{') and output.endswith('}'):
            data = json.loads(output)
            return {str(k): ';'.join(list(dict.fromkeys(str(v).split(';')))) for k, v in data.items()}
        elif isinstance(output, dict):
            return {str(k): ';'.join(list(dict.fromkeys(str(v).split(';')))) for k, v in output.items()}
    except json.JSONDecodeError as e:
        logger.warning(f"核心JSON解析失败: {str(e)}。 尝试备用方案。")

    # 步骤 4: 备用方案 - 逐行解析键值对
    result = {}
    # 兼容可能存在的 raw_output 场景
    if isinstance(output, str):
        lines = output.split('\n')
        key_pattern = re.compile(r'^\s*"?([^":]+)"?\s*[:：]\s*(.*)')
        for line in lines:
            line = line.strip().rstrip(',')  # 移除前后空格和末尾的逗号
            match = key_pattern.match(line)
            if match:
                key = match.group(1).strip()
                value = match.group(2).strip()
                # 移除值两端的引号
                if value.startswith('"') and value.endswith('"'):
                    value = value[1:-1]
                result[key] = ';'.join(list(dict.fromkeys(value.split(';'))))

    # 步骤 5: 最终备用方案
    if not result and output:
        logger.warning("所有解析方法均失败，返回原始输出。")
        return {"raw_output": str(output)}

    return result


def format_output(data: Dict[str, str], format_type: str) -> Optional[str]:
    """根据指定类型格式化提取出的数据。"""
    format_type = format_type.lower()
    try:
        if format_type == "json":
            return json.dumps(data, ensure_ascii=False, indent=2)
        elif format_type == "text":
            return "\n".join([f"{k}: {v}" for k, v in data.items() if v is not None])
        elif format_type == "xml":
            root = ET.Element("Result")
            for k, v in data.items():
                if v is None: continue
                safe_key = re.sub(r'[^\w]+', '_', k) if k else "item"
                if safe_key and safe_key[0].isdigit(): safe_key = "_" + safe_key
                if not safe_key: safe_key = "item"
                item = ET.SubElement(root, safe_key)
                item.text = str(v)
            from xml.dom import minidom
            return minidom.parseString(ET.tostring(root, 'unicode')).toprettyxml(indent="  ")
        elif format_type == "markdown":
            return "\n\n".join([f"**{k}**: {str(v) if v is not None else '无'}" for k, v in data.items()])
        else:
            return None
    except Exception as e:
        logger.error(f"格式化输出时出错 ({format_type}): {str(e)}", exc_info=True)
        return f"格式化为 {format_type} 时出错: {str(e)}"

def ext_api_call(payload):
    if EXT_API_URL is None or EXT_API_KEY is None:
        return {"code": 500, "msg": "EXT服务未配置", "data": None}
    headers = {'Authorization': f'Bearer {EXT_API_KEY}', 'Content-Type': 'application/json'}
    try:
        response = requests.post(EXT_API_URL, json=payload, headers=headers)
        response.raise_for_status()
        return response.json()
    except Exception as e:
        return {"code": 500, "msg": f"EXT API调用失败: {e}", "data": None}

async def main_test():
    """
    用于测试本地文件的异步主函数。
    """
    print("--- 开始本地文件抽取测试 ---")

    # 1. 指定您的本地文件路径
    #    请确保文件使用 UTF-8 编码
    file_path = r"C:\Users\19645\OneDrive\Desktop\full.md"  # <--- 修改这里

    # 2. 指定您要下达的抽取指令
    instruction = "商洛市仙娥湖片区道路桥梁工程项目共包含4条道路，4个子项的名称"  # <--- 修改这里

    # 3. 指定输出格式 (json, markdown, text, xml)
    output_format = "json"
    # ----------------------------------------------------

    try:
        print(f"正在读取文件: {file_path}")
        with open(file_path, 'r', encoding='utf-8') as f:
            text_content = f.read()
        print(f"文件读取成功，共 {len(text_content)} 个字符。")

        # 调用核心处理函数
        print("\n正在调用信息抽取模型，请稍候...")
        result = await process_content(
            text=text_content,
            instruction=instruction,
            output_format=output_format
        )

        # 格式化并打印结果
        print("\n--- 抽取结果 ---")
        if result and result.get("code") == 200:
            formatted_data = result.get("data", {}).get("formatted_data", "没有格式化数据。")
            print(formatted_data)
        else:
            print("处理失败，返回结果如下：")
            print(json.dumps(result, ensure_ascii=False, indent=2))

    except FileNotFoundError:
        print(f"\n错误：测试文件未找到！请确认路径 '{file_path}' 是否正确。")
    except Exception as e:
        print(f"\n测试过程中发生未知错误: {e}")
        traceback.print_exc()


if __name__ == '__main__':
    # 这个代码块允许您通过 `python ext.py` 直接运行测试

    # 1. 初始化服务
    print("正在为独立测试初始化服务配置...")
    config = ConfigParser()
    try:
        config.read('config.ini', encoding='utf-8')
        init_ext(config)
    except Exception as e:
        print(f"加载 config.ini 失败，测试无法继续: {e}")
        exit()

    # 2. 运行异步测试函数
    asyncio.run(main_test())